from ada_boost_utils import MyAdaBoost, data_from_csv
from matplotlib import pyplot as plt
from matplotlib.colors import ListedColormap
import numpy as np

# show data & prediction
def show_fig(data:np.ndarray, ada_boost:MyAdaBoost, width:int=500, height:int=500,):
    x, label = data[:, :-1], data[:, -1]
    x_positive, x_negative = data[label==1], data[label==-1]
    custom_cmap = ListedColormap(['#9CCC65', '#FFF176'])
    x_grid, y_grid = np.meshgrid(
        np.linspace(x[:, 0].min()-10, x[:, 0].max()+10, width),
        np.linspace(x[:, 1].min()-10, x[:, 1].max()+10, height),
    )
    x_mesh = np.c_[x_grid.ravel(), y_grid.ravel()]
    preds = ada_boost(x_mesh)
    preds = preds.reshape(x_grid.shape)     # its shape has to be the same as x_grid
    # get meshmap using the trained AdaBoost
    plt.figure()
    plt.pcolormesh(x_grid, y_grid, preds, cmap=custom_cmap, shading="auto")
    # plot raw data
    plt.scatter(x_positive[:, 0], x_positive[:, 1], marker='+', color='b')
    plt.scatter(x_negative[:, 0], x_negative[:, 1], marker='_', color='r')
    for i in range(data.shape[0]):
        plt.annotate(
            f"({ada_boost.weights[i]:.3f})", 
            xy=(x[i, 0], x[i, 1]), 
            xytext=(-10, -10), 
            textcoords='offset points'
        )
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title(f"AdaBoost Classifier(base classifiers={ada_boost.number_basis})")
    plt.xlim(x[:, 0].min()-10, x[:, 0].max()+10)
    plt.ylim(x[:, 1].min()-10, x[:, 1].max()+10)
    plt.show()
    pass

if __name__ == "__main__":
    ada_boost = MyAdaBoost('./data.csv', number_basis=3)
    ada_boost.train()
    accuracy = ada_boost.test(ada_boost.data)
    print(f"Accuracy: {accuracy*100:.2f}%")

    data = data_from_csv('./data.csv')
    show_fig(data, ada_boost)